import pickle
import pandas as pd
import datetime
import evaluation_indicators as ev
import numpy as np
import csv
from machine_learning import train_model
import math
from fileIO import FileIO
from dayCalculation import *
from log import Log
from Trader import Trader
from MLtrader import MLTrader


#回测系统类
class BackTest:
    def __init__(self, initial_money, benchmark_info, adjust_period, start_time, end_time,
                 stocks_info, trader, charge=0):
        self.initial_money = initial_money
        self.adjust_period = adjust_period
        self.benchmark_info = benchmark_info
        self.start_time = start_time
        self.end_time = end_time
        self.trader = trader
        self.log = Log()
        self.trade_days = list(stocks_info[list(stocks_info.keys())[0]].keys())

    def run(self):
        cur_time = self.start_time
        cur_money = self.initial_money
        benchmark_share = -1
        while isLater(cur_time, self.end_time):
            if cur_time not in self.trade_days:
                cur_time = next_day(cur_time, 1)
                continue
            # print(f'Day:{cur_time}')
            if benchmark_share == -1:
                benchmark_share = self.initial_money / self.benchmark_info[cur_time][1]
            self.trader.strategy(cur_time)
            benchmark_money = benchmark_share * self.benchmark_info[cur_time][1]
            benchmark_rate = benchmark_money / self.initial_money - 1
            earning_rate = self.trader.net_worth / self.initial_money - 1
            self.log.add_log(cur_time, self.trader.net_worth, earning_rate, benchmark_money, benchmark_rate)
            cur_time = next_day(cur_time, self.adjust_period)
        self.log.calculate()
        self.log.draw()
